"""


Version: 0.1
Author: lk
Date: 2022-03-08
"""
import warnings
import pprint
import numpy as np
import torch.nn.functional as F
import torch
from scipy import sparse
from torch.utils.data import DataLoader, Dataset
import dgl
from dgl import data as ddata
import networkx as nx
import matplotlib.pyplot as plt
import scipy
from scipy import linalg
from scipy.sparse import linalg

warnings.filterwarnings("ignore")
raw_dir = "."


class GraphBacher():
    def __init__(self, graph_q, graph_k):
        self.graph_q = graph_q
        self.graph_k = graph_k

    def __str__(self):
        return "graph_q is:" + str(self.graph_q) + "\n" + "graph_k is:" + str(self.graph_k)


def batcher(batch):
    graph_q, graph_k = zip(*batch)
    graph_q, graph_k = dgl.batch(graph_q), dgl.batch(graph_k)
    return GraphBacher(graph_q, graph_k)
def draw(g:dgl.DGLGraph):
    nx.draw(dgl.to_networkx(g))
    plt.show()


class GraphDataset(Dataset):
    def __init__(self, rw_hops=5,restart_prob=0.6,hidden_size=20):
        self.hidden_size = hidden_size
        self.restart_prob = restart_prob
        self.rw_hops = rw_hops
        graph_list = [
            ddata.CoauthorCSDataset(raw_dir=raw_dir),
            # ddata.CoauthorPhysicsDataset(raw_dir=raw_dir),
            # ddata.AmazonCoBuyPhotoDataset(raw_dir=raw_dir),
            # ddata.AmazonCoBuyComputerDataset(raw_dir=raw_dir)
        ]
        graph_list = list(map(lambda x: dgl.to_bidirected(x[0]), graph_list))
        # pprint.pprint(graph_list)
        self.graph = dgl.batch(graph_list, node_attrs=None, edge_attrs=None)
        self.graph.remove_nodes((self.graph.in_degrees()).nonzero().squeeze())

    def __len__(self):
        return self.graph.number_of_nodes()

    def get_2_traces(self, item):
        traces = []
        for i in range(2):
            traces.append(dgl.sampling.random_walk(
                g=self.graph,
                nodes=[item],
                length=self.rw_hops,
                # restart_prob=self.restart_prob

            )[0])
        return torch.cat(traces, dim=0)

    def __getitem__(self, item):
        traces = self.get_2_traces(item)
        # return traces[0],traces[1]
        q_nodes, k_nodes = traces[0], traces[1]
        graph_q,graph_k =dgl.node_subgraph(self.graph, q_nodes), dgl.node_subgraph(self.graph, k_nodes)
        graph_q = self.add_graph_features(graph_q)
        graph_k = self.add_graph_features(graph_k)
        return graph_q,graph_k
        # return dgl.node_subgraph(self.graph, q_nodes), dgl.node_subgraph(self.graph, k_nodes)
        # print(traces)
    def add_graph_features(self, g:dgl.DGLGraph):
        n=g.number_of_nodes()
        adj = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
        norm = sparse.diags(
            np.array(g.in_degrees().clip(1) ** -0.5),
            dtype=float
        )
        laplacian = norm * adj * norm
        # u, s, _ = sparse.linalg.svds(laplacian, k=min(n - 1, self.hidden_size), which='LM', return_singular_vectors='u')
        # res = scipy.linalg.svd(laplacian)
        u,s,_=sparse.linalg.svds(laplacian,k=5)
        x = u * sparse.diags(np.sqrt(s))
        x=torch.from_numpy(x)
        if n - 1 < self.hidden_size:
            x = F.pad(x, (0, self.hidden_size - n + 1), 'constant', 0)
        g.ndata['x'] = x
        # draw(g)
        return g


if __name__ == '__main__':
    torch.manual_seed(0)
    dataset = GraphDataset(
        rw_hops=5,
        restart_prob=0.8
    )
    dataloader = DataLoader(
        dataset,
        batch_size=128,
        shuffle=True,
        collate_fn=batcher,
        num_workers=4
    )
    for batch in dataloader:
        print(batch)
        # break
